import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import mnist
from tensorflow.keras.callbacks import LearningRateScheduler
import numpy as np

# 定义学习率调度函数
def lr_schedule(epoch):
    """
    学习率调度函数：每个epoch结束时将学习率衰减为初始学习率的0.9倍
    """
    initial_lr = 0.001
    lr = initial_lr * np.power(0.9, epoch)
    return lr

# 加载 MNIST 数据集
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
# 对数据进行预处理
train_images = train_images.reshape((60000, 28, 28, 1)).astype('float32') / 255
test_images = test_images.reshape((10000, 28, 28, 1)).astype('float32') / 255
# 将标签进行 One-hot 编码
train_labels = tf.keras.utils.to_categorical(train_labels)
test_labels = tf.keras.utils.to_categorical(test_labels)
# 构建卷积神经网络模型
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))
# 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 定义学习率调度器
lr_scheduler = LearningRateScheduler(lr_schedule)

# 训练模型
history = model.fit(train_images, train_labels, 
                    epochs=20, 
                    batch_size=64, 
                    validation_split=0.2, 
                    callbacks=[lr_scheduler])

# 评估模型
test_loss, test_acc = model.evaluate(test_images, test_labels)
print('Test accuracy:', test_acc)
# 保存模型
model.save('mnist_cnn_model.h5')
